#!/usr/bin/env bash

float_model_root="$HOME/onnx-model/public-model-zoo"
quant_model_root="$HOME/onnx-quants-3.8/ws-dev/onnx-quantization/onnx_quantize_tool/tmp"
zoo_root="/data/public_model_zoo/image/classification"
# zoo_root="$HOME/tmp/test"

function usage() {
  echo "Usage: $( basename $0 ) [ -m <model_name> | -h ]"
  echo -e "\t-m <model_name>: 要复制的模型主名，缺省复制${model_root}下所有模型"
  echo -e "\t-h: 输出本信息"
  exit 0
}

while getopts "hm:" arg
do
    case $arg in
        m)
          targetModel=${OPTARG}
          ;;
        *)
          usage
          ;;
    esac
done

function copy_float_model() {
  modelName=$1
  srcFile=$2

  fModelRoot="${zoo_root}/float"
  if [ ! -d ${fModelRoot} ]
  then
    mkdir -p "${fModelRoot}"
  fi

  if [ ! -f "${srcFile}" ]
  then
    echo "Error, model not found ${srcFile}"
    exit 2
  fi

  dstDir="${fModelRoot}/${modelName}"
  if [ ! -d "${dstDir}" ]
  then
    mkdir -v "${dstDir}"
  fi

  cp -v "${srcFile}" "${dstDir}"
}

function copy_quant_model() {
  chip=$1
  modelName=$2
  srcFile=$3

  qModelRoot="${zoo_root}/quant"
  if [ ! -d ${qModelRoot} ]
  then
    mkdir -p "${fModelRoot}"
  fi

  src="${quant_model_root}/${chip}/${modelName}_quantize.onnx"
  if [ ! -f "${src}" ]
  then
    echo "model not found ${src}"
    return
  fi

  dstDir="${qModelRoot}/${modelName}/${chip}"
  if [ ! -d "${dstDir}" ]
  then
    mkdir -v -p "${dstDir}"
  fi

  cp -v "${src}" "${dstDir}"
}

for floatModelPath in "${float_model_root}"/*
do
  modelStem="$( basename -s .onnx ${floatModelPath} )"
  if [ -n "${targetModel}" -a "${targetModel}" != "${modelStem}" ]
  then
    continue
  fi
  echo "------------------- ${modelStem} -------------------"
  copy_float_model "${modelStem}" "${floatModelPath}"
  for chip in TX510 TX511 DT53 DT56 DT57
  do
    copy_quant_model "${chip}" "${modelStem}" "${floatModelPath}"
  done
  echo "*********************** Done ***********************"
done
